"""BLE Data Transfer Service with fragmentation support.

Provides GATT characteristics for reliable data transfer
using the fragmentation protocol.
"""

import asyncio
import logging
import json
import struct
from typing import Optional, Callable, Dict, Any
from enum import Enum

from ..fragmentation import (
    FragmentationProtocol,
    TransferState,
    Frame
)

logger = logging.getLogger(__name__)


# Service and Characteristic UUIDs
DATA_TRANSFER_SERVICE_UUID = "0000FF30-0000-1000-8000-00805F9B34FB"
DATA_TX_UUID = "0000FF32-0000-1000-8000-00805F9B34FB"  # Write
DATA_RX_UUID = "0000FF33-0000-1000-8000-00805F9B34FB"  # Read, Notify
CONTROL_UUID = "0000FF34-0000-1000-8000-00805F9B34FB"  # Write
STATUS_UUID = "0000FF35-0000-1000-8000-00805F9B34FB"   # Read, Notify


class ControlCommand(Enum):
    """Transfer control commands."""
    START = 0x01
    STOP = 0x02
    CANCEL = 0x03
    RESET = 0x04
    GET_STATUS = 0x05


class DataTransferService:
    """BLE Data Transfer Service implementation."""
    
    def __init__(self, ack_required: bool = True,
                 max_retries: int = 3,
                 reassembly_timeout: float = 30.0):
        """Initialize data transfer service.
        
        Args:
            ack_required: Whether acknowledgments are required
            max_retries: Maximum retransmission attempts
            reassembly_timeout: Timeout for reassembly
        """
        # Initialize fragmentation protocol
        self.protocol = FragmentationProtocol(
            ack_required=ack_required,
            max_retries=max_retries,
            reassembly_timeout=reassembly_timeout
        )
        
        # Callbacks for BLE communication
        self.notify_callback: Optional[Callable[[str, bytes], None]] = None
        
        # Data buffers
        self.rx_buffer = asyncio.Queue(maxsize=100)
        self.tx_buffer = asyncio.Queue(maxsize=100)
        
        # Transfer state
        self.current_transfer: Optional[bytes] = None
        self.transfer_complete = asyncio.Event()
        
        # Setup protocol callbacks
        self.protocol.set_callbacks(
            send_callback=self._on_frame_send,
            receive_callback=self._on_data_complete
        )
        self.protocol.set_progress_callback(self._on_progress_update)
        
        # Service registration data
        self.service_data = self._create_service_definition()
        
        logger.info("Data Transfer Service initialized")
    
    def _create_service_definition(self) -> Dict[str, Any]:
        """Create GATT service definition.
        
        Returns:
            Service definition dictionary
        """
        return {
            "uuid": DATA_TRANSFER_SERVICE_UUID,
            "primary": True,
            "characteristics": [
                {
                    "uuid": DATA_TX_UUID,
                    "flags": ["write", "write-without-response"],
                    "value": bytearray(20),
                    "descriptors": []
                },
                {
                    "uuid": DATA_RX_UUID,
                    "flags": ["read", "notify"],
                    "value": bytearray(20),
                    "descriptors": [
                        {
                            "uuid": "00002902-0000-1000-8000-00805F9B34FB",
                            "flags": ["read", "write"],
                            "value": bytearray([0x00, 0x00])
                        }
                    ]
                },
                {
                    "uuid": CONTROL_UUID,
                    "flags": ["write"],
                    "value": bytearray(1),
                    "descriptors": []
                },
                {
                    "uuid": STATUS_UUID,
                    "flags": ["read", "notify"],
                    "value": bytearray(12),
                    "descriptors": [
                        {
                            "uuid": "00002902-0000-1000-8000-00805F9B34FB",
                            "flags": ["read", "write"],
                            "value": bytearray([0x00, 0x00])
                        }
                    ]
                }
            ]
        }
    
    async def handle_write_tx(self, value: bytes) -> bool:
        """Handle write to DATA_TX characteristic.
        
        Args:
            value: Written value (frame data)
            
        Returns:
            True if handled successfully
        """
        try:
            # Receive frame through protocol
            self.protocol.receive(value)
            return True
            
        except Exception as e:
            logger.error(f"Error handling TX write: {e}")
            return False
    
    async def handle_read_rx(self) -> bytes:
        """Handle read from DATA_RX characteristic.
        
        Returns:
            Next frame data or empty if none
        """
        try:
            # Get next frame from TX buffer (non-blocking)
            if not self.tx_buffer.empty():
                frame_data = self.tx_buffer.get_nowait()
                return frame_data
            return b'\x00' * 20
            
        except Exception as e:
            logger.error(f"Error handling RX read: {e}")
            return b'\x00' * 20
    
    async def handle_write_control(self, value: bytes) -> bool:
        """Handle write to CONTROL characteristic.
        
        Args:
            value: Control command
            
        Returns:
            True if handled successfully
        """
        try:
            if not value:
                return False
            
            command = ControlCommand(value[0])
            
            if command == ControlCommand.START:
                # Start transfer
                logger.info("Transfer started")
                self.transfer_complete.clear()
                
            elif command == ControlCommand.STOP:
                # Stop transfer
                logger.info("Transfer stopped")
                self.protocol.cancel()
                
            elif command == ControlCommand.CANCEL:
                # Cancel transfer
                logger.info("Transfer cancelled")
                self.protocol.cancel()
                self.transfer_complete.set()
                
            elif command == ControlCommand.RESET:
                # Reset protocol
                logger.info("Protocol reset")
                self.protocol.reset()
                await self._clear_buffers()
                
            elif command == ControlCommand.GET_STATUS:
                # Trigger status update
                await self._send_status_update()
            
            return True
            
        except Exception as e:
            logger.error(f"Error handling control command: {e}")
            return False
    
    async def handle_read_status(self) -> bytes:
        """Handle read from STATUS characteristic.
        
        Returns:
            Status data (12 bytes)
        """
        try:
            status = self.protocol.get_status()
            
            # Pack status into 12 bytes
            # Format: state(1) + progress(1) + transferred(4) + throughput(2) + errors(2) + reserved(2)
            state_map = {
                "IDLE": 0,
                "SENDING": 1,
                "RECEIVING": 2,
                "WAITING_ACK": 3,
                "RETRANSMITTING": 4,
                "COMPLETE": 5,
                "ERROR": 6,
                "CANCELLED": 7
            }
            
            state = state_map.get(status["state"], 0)
            progress = min(100, status["progress"])
            transferred = min(0xFFFFFFFF, status["transferred"])
            throughput = min(0xFFFF, int(status["throughput"]))
            errors = min(0xFFFF, int(status["error_rate"] * 100))
            
            return struct.pack(
                '>BBIHH2x',
                state,
                progress,
                transferred,
                throughput,
                errors
            )
            
        except Exception as e:
            logger.error(f"Error reading status: {e}")
            return b'\x00' * 12
    
    async def send_data(self, data: bytes) -> bool:
        """Send data using fragmentation protocol.
        
        Args:
            data: Data to send
            
        Returns:
            True if successful
        """
        try:
            self.current_transfer = data
            self.transfer_complete.clear()
            
            # Send through protocol
            success = await self.protocol.send(data)
            
            if success:
                logger.info(f"Data sent successfully: {len(data)} bytes")
            else:
                logger.error("Failed to send data")
            
            return success
            
        except Exception as e:
            logger.error(f"Error sending data: {e}")
            return False
    
    async def receive_data(self, timeout: float = 30.0) -> Optional[bytes]:
        """Wait for complete data reception.
        
        Args:
            timeout: Reception timeout
            
        Returns:
            Received data or None if timeout
        """
        try:
            # Wait for data from RX buffer
            data = await asyncio.wait_for(
                self.rx_buffer.get(),
                timeout=timeout
            )
            return data
            
        except asyncio.TimeoutError:
            logger.warning("Receive timeout")
            return None
        except Exception as e:
            logger.error(f"Error receiving data: {e}")
            return None
    
    def _on_frame_send(self, frame_data: bytes):
        """Callback when frame needs to be sent.
        
        Args:
            frame_data: Frame data to send
        """
        try:
            # Add to TX buffer
            self.tx_buffer.put_nowait(frame_data)
            
            # Notify if callback is set
            if self.notify_callback:
                self.notify_callback(DATA_RX_UUID, frame_data)
            
        except asyncio.QueueFull:
            logger.warning("TX buffer full, dropping frame")
        except Exception as e:
            logger.error(f"Error in frame send callback: {e}")
    
    def _on_data_complete(self, data: bytes):
        """Callback when complete data is received.
        
        Args:
            data: Complete received data
        """
        try:
            # Add to RX buffer
            self.rx_buffer.put_nowait(data)
            
            # Signal completion
            self.transfer_complete.set()
            
            logger.info(f"Data reception complete: {len(data)} bytes")
            
        except asyncio.QueueFull:
            logger.warning("RX buffer full, dropping data")
        except Exception as e:
            logger.error(f"Error in data complete callback: {e}")
    
    def _on_progress_update(self, progress: float):
        """Callback for transfer progress updates.
        
        Args:
            progress: Progress percentage (0-100)
        """
        logger.debug(f"Transfer progress: {progress:.1f}%")
        
        # Send status notification
        asyncio.create_task(self._send_status_update())
    
    async def _send_status_update(self):
        """Send status update notification."""
        try:
            status_data = await self.handle_read_status()
            
            if self.notify_callback:
                self.notify_callback(STATUS_UUID, status_data)
                
        except Exception as e:
            logger.error(f"Error sending status update: {e}")
    
    async def _clear_buffers(self):
        """Clear all data buffers."""
        # Clear RX buffer
        while not self.rx_buffer.empty():
            try:
                self.rx_buffer.get_nowait()
            except asyncio.QueueEmpty:
                break
        
        # Clear TX buffer
        while not self.tx_buffer.empty():
            try:
                self.tx_buffer.get_nowait()
            except asyncio.QueueEmpty:
                break
        
        logger.debug("Buffers cleared")
    
    def set_notify_callback(self, callback: Callable[[str, bytes], None]):
        """Set notification callback.
        
        Args:
            callback: Callback for notifications (uuid, data)
        """
        self.notify_callback = callback
    
    def get_service_definition(self) -> Dict[str, Any]:
        """Get GATT service definition.
        
        Returns:
            Service definition for registration
        """
        return self.service_data
    
    async def cleanup(self):
        """Clean up service resources."""
        try:
            # Cancel transfer
            self.protocol.cancel()
            
            # Clear buffers
            await self._clear_buffers()
            
            logger.info("Data Transfer Service cleaned up")
            
        except Exception as e:
            logger.error(f"Cleanup error: {e}")


# Integration helper
class DataTransferIntegration:
    """Helper for integrating data transfer service with GATT server."""
    
    def __init__(self, gatt_server):
        """Initialize integration.
        
        Args:
            gatt_server: GATT server instance
        """
        self.gatt_server = gatt_server
        self.service = DataTransferService()
        
        # Set notification callback
        self.service.set_notify_callback(self._notify_clients)
    
    def register_service(self):
        """Register data transfer service with GATT server."""
        service_def = self.service.get_service_definition()
        
        # Register with GATT server
        # This would be implementation-specific based on your GATT server
        logger.info("Data Transfer Service registered with GATT server")
    
    def _notify_clients(self, uuid: str, data: bytes):
        """Send notification to connected clients.
        
        Args:
            uuid: Characteristic UUID
            data: Notification data
        """
        # Implementation depends on GATT server
        # This would send notification to all subscribed clients
        logger.debug(f"Notifying clients on {uuid}: {len(data)} bytes")
    
    async def handle_characteristic_write(self, uuid: str, value: bytes) -> bool:
        """Route characteristic writes to service.
        
        Args:
            uuid: Characteristic UUID
            value: Written value
            
        Returns:
            True if handled
        """
        if uuid == DATA_TX_UUID:
            return await self.service.handle_write_tx(value)
        elif uuid == CONTROL_UUID:
            return await self.service.handle_write_control(value)
        
        return False
    
    async def handle_characteristic_read(self, uuid: str) -> bytes:
        """Route characteristic reads to service.
        
        Args:
            uuid: Characteristic UUID
            
        Returns:
            Characteristic value
        """
        if uuid == DATA_RX_UUID:
            return await self.service.handle_read_rx()
        elif uuid == STATUS_UUID:
            return await self.service.handle_read_status()
        
        return b''


# Example usage
if __name__ == "__main__":
    import asyncio
    
    async def test_service():
        """Test data transfer service."""
        # Create service
        service = DataTransferService()
        
        # Test data
        test_data = b"Hello, BLE Data Transfer Service!" * 5
        
        print(f"Testing with {len(test_data)} bytes of data")
        
        # Simulate sending data
        send_task = asyncio.create_task(service.send_data(test_data))
        
        # Simulate receiving frames and sending ACKs
        async def simulate_ble_transfer():
            """Simulate BLE frame transfer."""
            while True:
                # Get frame from TX buffer
                try:
                    frame_data = await asyncio.wait_for(
                        service.handle_read_rx(),
                        timeout=0.1
                    )
                    
                    if frame_data != b'\x00' * 20:
                        print(f"Frame sent via BLE: {len(frame_data)} bytes")
                        
                        # Echo back to RX (simulating loopback)
                        await service.handle_write_tx(frame_data)
                        
                except asyncio.TimeoutError:
                    continue
                except Exception as e:
                    print(f"Error in simulation: {e}")
                    break
        
        # Start simulation
        sim_task = asyncio.create_task(simulate_ble_transfer())
        
        # Wait for send to complete
        success = await send_task
        print(f"Send {'successful' if success else 'failed'}")
        
        # Try to receive data
        received = await service.receive_data(timeout=5.0)
        if received:
            print(f"Received {len(received)} bytes")
            print(f"Data matches: {received == test_data}")
        else:
            print("No data received")
        
        # Get status
        status_data = await service.handle_read_status()
        status = struct.unpack('>BBIHH2x', status_data)
        print(f"Status: state={status[0]}, progress={status[1]}%, "
              f"transferred={status[2]}, throughput={status[3]} B/s")
        
        # Cleanup
        sim_task.cancel()
        await service.cleanup()
    
    # Run test
    asyncio.run(test_service())